#include "mex.h"
#include <stdint.h>
#include <stdio.h>

typedef struct node {
    struct node* next;
    uint32_t sdst;
    mwSize nn;
} node;

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
    /* Declare variables. */
    mwSize mX, mY, N, i, j, t, *dims;
    uint8_t * pX, *pY;
    uint32_t val, sum;
    uint32_t *xx, *yy, *dd;
    mwSize* nn;
    node *snode, *enode, *tnode, *pnode;
    double* pk;
    int num, k;

    /* Check for proper number of input and output arguments. */
    if (nrhs != 2 && nrhs!=3) {
        mexErrMsgTxt("Two or three input arguments required.");
    }
    if (nlhs > 3) {
        mexErrMsgTxt("Too many output arguments.");
    }


    /* Check data type of input argument. */
    if (!(mxIsUint8(prhs[0]))) {
        mexErrMsgTxt("Input array must be of type uint8.");
    }

    /* Get the number of elements in the input argument. */
    mX = mxGetM(prhs[0]);
    N = mxGetN(prhs[0]);

    /* Get the data. */
    pX  = (uint8_t *) mxGetData(prhs[0]);

    xx = malloc(mX*sizeof(uint32_t));
    for (j = 0; j < mX; ++j) {
        xx[j] = 0;
        for (t = 0; t < N; ++t) {
            if(pX[j+mX*t] != 0)
                xx[j] += pX[j+mX*t]*pX[j+mX*t];
        }
    }

    if (nrhs == 2) {
        pk  = (double*) mxGetData(prhs[1]);
        pY = pX;
        mY = mX;
        yy = xx;
    } else {
        if (N != mxGetN(prhs[1]))
            mexErrMsgTxt("Dimension mismatch!");
        pY = (uint8_t *) mxGetData(prhs[1]);
        mY = mxGetM(prhs[1]);
        yy = malloc(mY*sizeof(uint32_t));
        for (i = 0; i < mY; ++i) {
            yy[i] = 0;
            for (t = 0; t < N; ++t) {
                if(pY[j+mY*t] != 0)
                    yy[i] += pY[j+mY*t]*pY[j+mY*t];
            }
        }
        pk  = (double *)mxGetPr(prhs[1]);
    }
    num = (int) pk[0];

    //printf("pk = %f\n", *pk);

    snode = (node*) malloc(sizeof(node));
    snode->next = NULL;
    snode->sdst = 0;
    enode = snode;
    for (k = 0; k < num; ++k) {
        tnode = enode;
        enode = (node*) malloc(sizeof(node));
        enode->sdst = UINT32_MAX;
        enode->next = tnode;
    }

    //printf("Data loaded.\n");

    dims = (mwSize*) malloc(2*sizeof(mwSize));
    dims[0] = mY;
    dims[1] = num;

    plhs[0] = mxCreateNumericArray(2, dims, mxUINT32_CLASS, mxREAL);
    plhs[1] = mxCreateNumericArray(2, dims, mxUINT32_CLASS, mxREAL);

    nn = mxGetData(plhs[0]);
    dd = mxGetData(plhs[1]);

    //printf("Output set.\n");

    for (i = 0; i < mY; ++i) {
        tnode = enode;
        for (k = 0; k < num; ++k) {
            tnode->sdst = UINT32_MAX;
            tnode = tnode->next;
        }
        for (j = 0; j < mX; ++j) {
            //printf("i = %d; j = %d\n", i, j);
            sum = 0;
            for (t = 0; t < N; ++t) {
                if ((pY[i+mY*t]!=0)&&(pX[j+mX*t]!=0)) {
                    sum += pY[i+mY*t]*pX[j+mX*t];
                }
            }
            val = xx[j] + yy[i] - 2*sum;
            if (val < enode->sdst) {
                tnode = enode;
                while (val < tnode->next->sdst) tnode = tnode->next;
                //if (i == 0) printf("J = %lu; Value = %lu; Next = %lu\n", j, val, tnode->sdst);
                if (tnode == enode) {
                    enode->sdst = val;
                    enode->nn = j;
                } else {
                    pnode = enode;
                    enode = enode->next;
                    pnode->nn = j;
                    pnode->sdst = val;
                    pnode->next = tnode->next;
                    tnode->next = pnode;
                }
            }

            //printf("NUM = %d\n", num);


        }
        //printf("nnM = %d; nnN = %d\n", mxGetM(plhs[0]), mxGetN(plhs[0]));
        tnode = enode;
        for (k = num-1; k >= 0; --k) {
            dd[i+mY*k] = tnode->sdst;
            nn[i+mY*k] = tnode->nn+1;
            tnode = tnode->next;
        }
    }
}
